import numpy as np
from robot_arm import RobotArm
from optimizer import PathOptimizer
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

def solve_problem1():
    # 创建机器人臂实例
    robot = RobotArm()
    optimizer = PathOptimizer(robot)
    
    # 目标点位置
    target_position = np.array([1500, 1200, 200])
    
    # 优化路径
    optimal_angles, final_error = optimizer.optimize_path(
        target_position,
        error_weight=1.0,
        energy_weight=0.0
    )
    
    # 获取最终位置
    final_position = robot.calculate_end_effector_position(optimal_angles)
    
    # 打印结果
    print("最优关节角度 (弧度):", optimal_angles)
    print("最优关节角度 (度):", np.rad2deg(optimal_angles))
    print("最终位置:", final_position)
    print("末端误差 (mm):", final_error)
    
    # 绘制机器人臂的零位状态图
    plot_robot_arm(robot)

def plot_robot_arm(robot: RobotArm):
    """绘制机器人臂的零位状态图"""
    fig = plt.figure(figsize=(10, 8))
    ax = fig.add_subplot(111, projection='3d')
    
    # 计算各个关节的位置
    T = np.eye(4)
    positions = [(0, 0, 0)]
    
    for i in range(6):
        theta = robot.dh_params['theta'][i]
        d = robot.dh_params['d'][i]
        a = robot.dh_params['a'][i]
        alpha = robot.dh_params['alpha'][i]
        
        T = T @ robot.transform_matrix(theta, d, a, alpha)
        positions.append(T[:3, 3])
    
    # 绘制连杆
    positions = np.array(positions)
    ax.plot(positions[:, 0], positions[:, 1], positions[:, 2], 'bo-', linewidth=2)
    
    # 设置坐标轴
    ax.set_xlabel('X (mm)')
    ax.set_ylabel('Y (mm)')
    ax.set_zlabel('Z (mm)')
    ax.set_title('六自由度机器臂零位状态图')
    
    # 设置视角
    ax.view_init(elev=20, azim=45)
    
    plt.show()

if __name__ == "__main__":
    solve_problem1() 